'''
2.并行化（Parallelization） 定义：将任务拆分为多个子任务，同时（并行）执行，以提升效率或获取多样性结果。

应用场景：
多角度分析：对同一段文本并行生成“正面”“负面”“中立”三种观点。
多模型投票（Ensemble）：多个模型同时处理同一任务，综合结果提升鲁棒性。
大规模数据处理：如批量处理1000条用户评论的情感分析，分批并行执行。
A/B测试生成：并行生成多个版本的广告文案，用于后续测试。
✅ 优势：提升吞吐量、支持多样性输出、增强系统容错性。
'''

from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from utils import init_llm, set_env, visualize_graph

# 初始化图，预定义存储的字段内容
class State(TypedDict):
    topic: str
    joke: str
    story: str
    poem: str
    combined_output: str


def call_llm_1(state: State):
    """第一个 LLM 调用生成初始笑话"""

    msg = llm.invoke(f"写一个关于 {state['topic']} 话题的笑话/no_think")
    return {"joke": msg.content}


def call_llm_2(state: State):
    """第二次 LLM 调用生成故事"""

    msg = llm.invoke(f"写一个关于 {state['topic']} 话题的故事/no_think")
    return {"story": msg.content}


def call_llm_3(state: State):
    """第三次 LLM 调用生成诗歌"""

    msg = llm.invoke(f"写一个关于 {state['topic']} 话题的诗歌/no_think")
    return {"poem": msg.content}

def aggregator(state: State):
    """将笑话和故事组合成一个输出"""

    combined = f"以下是关于 {state['topic']} 话题的故事、笑话、诗歌!\n\n"
    combined += f"故事:\n{state['story']}\n\n"
    combined += f"笑话:\n{state['joke']}\n\n"
    combined += f"诗歌:\n{state['poem']}"
    return {"combined_output": combined}


def build_graph():
    workflow = StateGraph(State)

    workflow.add_node("call_llm_1", call_llm_1)
    workflow.add_node("call_llm_2", call_llm_2)
    workflow.add_node("call_llm_3", call_llm_3)
    workflow.add_node("aggregator", aggregator)

    workflow.add_edge(START, "call_llm_1")
    workflow.add_edge(START, "call_llm_2")
    workflow.add_edge(START, "call_llm_3")
    workflow.add_edge("call_llm_1", "aggregator")
    workflow.add_edge("call_llm_2", "aggregator")
    workflow.add_edge("call_llm_3", "aggregator")
    workflow.add_edge("aggregator", END)

    graph = workflow.compile()

    return graph


if __name__ == "__main__":
    set_env()  # 设置环境变量
    llm = init_llm(temperature=0)  # 初始化LLM，温度设为0以获得确定性输出
    graph = build_graph()  # 构建工作流图
 
    visualize_graph(graph)

    state = graph.invoke({"topic": "猫"})
    print(state["combined_output"])



    